Before use this notebook¶

This notebook is based on 'https://colab.research.google.com/github/roboflow/notebooks/blob/main/notebooks/train-yolov10-object-detection-on-custom-dataset.ipynb'

Some modifications were made to train custom data.

In [2]:
import os
HOME = os.getcwd()
In [3]:
!nvidia-smi
!lscpu
!free

Sat Feb  8 00:26:20 2025       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.120                Driver Version: 560.94         CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3060 Ti     On  |   00000000:07:00.0  On |                  N/A |
|  0%   32C    P8             17W /  200W |     895MiB /   8192MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A       815      C   /python3.10                                 N/A      |
+-----------------------------------------------------------------------------------------+
Architecture:             x86_64
  CPU op-mode(s):         32-bit, 64-bit
  Address sizes:          48 bits physical, 48 bits virtual
  Byte Order:             Little Endian
CPU(s):                   6
  On-line CPU(s) list:    0-5
Vendor ID:                AuthenticAMD
  Model name:             AMD Ryzen 5 5600 6-Core Processor
    CPU family:           25
    Model:                33
    Thread(s) per core:   1
    Core(s) per socket:   6
    Socket(s):            1
    Stepping:             2
    BogoMIPS:             6986.95
    Flags:                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge m
                          ca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall
                           nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep
                          _good nopl tsc_reliable nonstop_tsc cpuid extd_apicid 
                          pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcn
                          t aes xsave avx f16c rdrand hypervisor lahf_lm cmp_leg
                          acy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch
                           osvw topoext perfctr_core ssbd ibrs ibpb stibp vmmcal
                          l fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx
                           smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 x
                          saves clzero xsaveerptr arat npt nrip_save tsc_scale v
                          mcb_clean flushbyasid decodeassists pausefilter pfthre
                          shold v_vmsave_vmload umip vaes vpclmulqdq rdpid fsrm
Virtualization features:  
  Virtualization:         AMD-V
  Hypervisor vendor:      Microsoft
  Virtualization type:    full
Caches (sum of all):      
  L1d:                    192 KiB (6 instances)
  L1i:                    192 KiB (6 instances)
  L2:                     3 MiB (6 instances)
  L3:                     32 MiB (1 instance)
Vulnerabilities:          
  Gather data sampling:   Not affected
  Itlb multihit:          Not affected
  L1tf:                   Not affected
  Mds:                    Not affected
  Meltdown:               Not affected
  Mmio stale data:        Not affected
  Reg file data sampling: Not affected
  Retbleed:               Not affected
  Spec rstack overflow:   Mitigation; safe RET, no microcode
  Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prct
                          l and seccomp
  Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointe
                          r sanitization
  Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STI
                          BP disabled; RSB filling; PBRSB-eIBRS Not affected; BH
                          I Not affected
  Srbds:                  Not affected
  Tsx async abort:        Not affected
               total        used        free      shared  buff/cache   available
Mem:         8137156     1535784     6072172       82260      529200     6280860
Swap:        4194304           0     4194304

Install YOLOv10¶

NOTE: Currently, YOLOv10 does not have its own PyPI package. Therefore, we need to install the code from the source.

In [8]:
!pip install -q git+https://github.com/THU-MIG/yolov10.git

NOTE: We will also install two additional packages: roboflow to download the dataset from Roboflow Universe, which we will use to train our model, and supervision, which we will use for visualizing the results.

In [ ]:
!pip install supervision roboflow

Download pre-trained weights¶

NOTE: YOLOv10 provides weight files pre-trained on the COCO dataset in various sizes. Let's download them.

In [ ]:
!mkdir -p {HOME}/weights
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10n.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10s.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10m.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10b.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10x.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10l.pt
!ls -lh {HOME}/weights

Custom Training¶

In [ ]:
%cd {HOME}

# Before run train command, Print it
!echo yolo task=detect mode=train epochs=10 batch=32 plots=True model={HOME}/weights/yolov10n.pt data={HOME}/data/data.yaml
In [ ]:
# Run training!
!yolo task=detect mode=train epochs=10 batch=32 plots=True model={HOME}/weights/yolov10n.pt data={HOME}/data/data.yaml

NOTE: change to last train directory

In [ ]:
last_train_dir=!ls -d {HOME}/runs/detect/train* | sort -V | tail -n 1
if len(last_train_dir):
    print(last_train_dir[0])
    global TRAIN_DIR
    TRAIN_DIR=last_train_dir[0]
else:
    print("no train folder!")

NOTE: Let's display train result

In [6]:
from IPython.display import Image
In [15]:
Image(filename=f'{TRAIN_DIR}/confusion_matrix.png', width=600)
Out[15]:
No description has been provided for this image
In [16]:
Image(filename=f'{TRAIN_DIR}/results.png', width=600)
Out[16]:
No description has been provided for this image

Inference with Custom Model¶

NOTE: Let's start by loading our newly trained model.

In [5]:
from ultralytics import YOLO
import supervision as sv

dataset = sv.DetectionDataset.from_yolo(
    images_directory_path=f"{HOME}/data/images/valid",
    annotations_directory_path=f"{HOME}/data/labels/valid",
    data_yaml_path=f"{HOME}/data/data.yaml"
)

bounding_box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()

NOTE: Let's display images from our validation set and visualize the results.

In [9]:
model = YOLO(f'{TRAIN_DIR}/weights/best.pt')

for path, image, annotation in dataset:
    results = model(source=image, conf=0.5)[0]
    detections = sv.Detections.from_ultralytics(results)

    annotated_image = bounding_box_annotator.annotate(scene=image, detections=detections)
    annotated_image = label_annotator.annotate(scene=image, detections=detections)

    sv.plot_image(annotated_image)

0: 384x640 4 _flawless_amethysts, 1 _flawless_diamond, 3 _flawless_saphires, 3 _perfect_diamonds, 3 _perfect_emeralds, 5 _perfect_rubys, 3 _perfect_saphires, 1 _perfect_skull, 26.5ms
Speed: 2.5ms preprocess, 26.5ms inference, 1.8ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image 
0: 384x640 3 _flawless_amethysts, 2 _flawless_emeralds, 1 _flawless_ruby, 1 _flawless_saphire, 6 _flawless_skulls, 3 _perfect_amethysts, 2 _perfect_diamonds, 4 _perfect_emeralds, 2 _perfect_rubys, 2 _perfect_saphires, 8 _perfect_skulls, 35.9ms
Speed: 2.7ms preprocess, 35.9ms inference, 2.3ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image 
0: 384x640 85 _perfect_amethysts, 14 _perfect_skulls, 24.4ms
Speed: 2.7ms preprocess, 24.4ms inference, 2.0ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image 
0: 384x640 2 _flawless_diamonds, 2 _flawless_emeralds, 3 _flawless_rubys, 1 _flawless_saphire, 1 _flawless_skull, 4 _perfect_diamonds, 23 _perfect_rubys, 1 _perfect_saphire, 25.8ms
Speed: 2.6ms preprocess, 25.8ms inference, 1.9ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image